import logging
from pathlib import Path
from typing import Any, Dict, Optional
import streamlit as st
import json

class StateManager:
    """状态管理器"""
    def __init__(self, session_dir: str = "chat_sessions"):
        self.session_dir = Path(session_dir)
        self.session_dir.mkdir(parents=True, exist_ok=True)
        self.logger = logging.getLogger(__name__)
        self._initialize_states()
    
    def _initialize_states(self):
        """初始化所有状态"""
        if "messages" not in st.session_state:
            st.session_state.messages = []
        if "model" not in st.session_state:
            st.session_state.model = "llama2"
        if "current_session" not in st.session_state:
            st.session_state.current_session = None
    
    def get_state(self, key: str, default: Any = None) -> Any:
        """获取状态值"""
        return getattr(st.session_state, key, default)
    
    def set_state(self, key: str, value: Any):
        """设置状态值"""
        setattr(st.session_state, key, value)
    
    def append_message(self, message: Dict):
        """添加消息"""
        st.session_state.messages.append(message)
        self._save_current_session()
    
    def _save_current_session(self):
        """保存当前会话"""
        if st.session_state.current_session:
            session_file = self.session_dir / f"{st.session_state.current_session}.json"
            try:
                with open(session_file, 'w', encoding='utf-8') as f:
                    json.dump(st.session_state.messages, f, ensure_ascii=False, indent=2)
            except Exception as e:
                self.logger.error(f"保存会话失败: {str(e)}") 